#ifndef DEVICE_QUERY_H
#define DEVICE_QUERY_H

#include <cuda_runtime.h>
#include <stdlib.h>
#include <assert.h>
#include <stdio.h>

#ifdef NVIDIA_GPU_ARCH
  constexpr unsigned kWarpSize = 32;
#elif defined(METAX_GPU_ARCH)
  constexpr unsigned kWarpSize = 64;
#else
  constexpr unsigned kWarpSize = 64;
#endif

#define checkGPUErrors(val) __check((val), #val, __FILE__, __LINE__)

void __check(cudaError_t result, char const *const func, const char *const file,
           int const line) {
  if (result) {
    fprintf(stderr, "CUDA error at %s:%d code=%d(%s) \"%s\" \n", file, line,
            static_cast<unsigned int>(result), cudaGetErrorName(result), func);
    exit(EXIT_FAILURE);
  }
}

#define checkGPULastErrors() __checkLastError(__FILE__, __LINE__)

inline void __checkLastError(const char *const file, int const line) {
  cudaError_t err = cudaGetLastError();
  if (cudaSuccess != err) {
    fprintf(stderr, "CUDA error at %s:%d code=%d(%s) \n", file, line,
            static_cast<unsigned int>(err), cudaGetErrorName(err));
    exit(EXIT_FAILURE);
  }
}

struct DeviceParams {
  unsigned NumSM;           // number of SMs
  unsigned WarpSize;           // max threads per warp
  unsigned MaxThreadsPerSM;  // max threads / sm
  unsigned MaxSharedMem; // Max configerable shared memory size in bytes
  unsigned MaxWarpsPerSM;    // max warps / sm
  unsigned MaxRegPerSM;      // max registers / sm

  unsigned MaxThreadsPerBlock;         // max threads per threadblock
  unsigned MaxSharedPerBlock; // Max configerable shared memory size
                                          // per block in bytes
  unsigned
      MaxRegPerBlock; // Max configerable shared memory size per block in bytes
  unsigned CoreFrequency; // SM clock freq in MHZ
  size_t L2Capacity; // L2 size in bytes

  size_t MemCapacity;            // Memory size in bytes
  unsigned MemFrequency; // Memory clock freq in MHZ
  unsigned MemBitWidth;      // Memory bit width
} gDeviceParams;

void initializeDeviceProp(unsigned deviceID) {
  cudaDeviceProp deviceProp;
  cudaSetDevice(deviceID);
  cudaGetDeviceProperties(&deviceProp, deviceID);

  // core stats
  gDeviceParams.NumSM =  deviceProp.multiProcessorCount;
  gDeviceParams.MaxThreadsPerSM =  deviceProp.maxThreadsPerMultiProcessor;
  gDeviceParams.MaxSharedMem =  deviceProp.sharedMemPerMultiprocessor;
  gDeviceParams.WarpSize =  deviceProp.warpSize;
  gDeviceParams.MaxWarpsPerSM = 
      deviceProp.maxThreadsPerMultiProcessor / deviceProp.warpSize;
  gDeviceParams.MaxRegPerSM =  deviceProp.regsPerMultiprocessor;
#if defined(METAX_GPU_ARCH)
  gDeviceParams.CoreFrequency = 1125;
#elif defined(NVIDIA_GPU_ARCH)
  gDeviceParams.CoreFrequency =  deviceProp.clockRate * 1e-3f;
#endif

  // threadblock stats
  gDeviceParams.MaxThreadsPerBlock =  deviceProp.maxThreadsPerBlock;
  gDeviceParams.MaxSharedPerBlock =  deviceProp.sharedMemPerBlock;
  gDeviceParams.MaxRegPerBlock =  deviceProp.regsPerBlock;

  // L2 cache
  gDeviceParams.L2Capacity =  deviceProp.l2CacheSize;

  // memory
  gDeviceParams.MemCapacity =  deviceProp.totalGlobalMem;
  gDeviceParams.MemFrequency =  deviceProp.memoryClockRate * 1e-3f;
  gDeviceParams.MemBitWidth =  deviceProp.memoryBusWidth;
}

void print_gpu_info(int deviceID) {
  cudaDeviceProp deviceProp;
  int devID = 0;
  cudaSetDevice(devID);
  auto error = cudaGetDeviceProperties(&deviceProp, devID);
  if (error != cudaSuccess) {
    printf("cudaGetDeviceProperties returned error code %d, line(%d)\n", error,
           __LINE__);
    exit(EXIT_FAILURE);
  }
  printf("GPU Device %d: \"%s\" with compute capability %d.%d\n\n", devID,
         deviceProp.name, deviceProp.major, deviceProp.minor);
}


#endif

